#
#   HarmonicOscillator1D.py
#
#   A Library to define the system's parameters, and its wavefunction.
#
#   It might be faster to write this in C or Fortran, because it gets
#   called over and over... but this is fast enough for our uses.
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from math import sqrt
from cmath import pi
from cmath import exp
from math import factorial

# Now define physics

hbar = 1.0
mass = 1.0

all_omega = [1.5,2.0] # number of dimensions determined dynamically from number of coefficients

def V(r):
    temp = map(lambda a,b: (a*b)**2, all_omega, r)
    return 0.5 * mass * sum(temp)

# Define a few eigenstates in each direction...

def Hermite_Coeff(n,l): # Nth polynomial, L=(N-2*M)th degree in x.
    m = float(n-l)/2.0
    try:
        return factorial(n)*(-1)**m/(factorial(m)*factorial(l))*2**l
    except ValueError: # if m is negative, fractional... then the coefficient is zero.
        return 0

def Hermite_Poly(n,omega):
    returnable_string = str(int(Hermite_Coeff(n,n))) + " * ("+str(sqrt(mass*omega/hbar))+"*x)**" + str(n)
    for i in range(n-2,-1,-2):
        returnable_string += " + " + str(int(Hermite_Coeff(n,i))) + " * ("+str(sqrt(mass*omega/hbar))+"*x)**" + str(i)
    return "("+returnable_string+")"

def make_eigenfunction_n(n,omega):
    # This generates the definition code for the Nth eigenfunction, as a string.
    # We try to evaluate as much of the algebra as possible at `compile time', because
    # profiling had revealed that the majority of runtime was spent in these eigenfunctions.
    
    normalisation = str(eval("1.0/sqrt(2.0**"+str(n)+" * factorial("+str(n)+"))*("+str((mass*omega/(hbar*pi)))+")**0.25"))
    gauss = str(exp(-mass*omega/hbar/2.0).real)+"**(x*x)"
    time = "exp(("+str((0-1j)*omega*(n+0.5))+")*t)"
    
    # Now, put everything together:
    
    # Given strings "a","b","c",... this function returns the string "a*b*c*..."
    str_product = lambda a,b,*args: str_product(a+"*"+b,*args) if args else a+"*"+b
    
    defn = "lambda t,x : " + str_product(normalisation, gauss, Hermite_Poly(n,omega), time)
    
    return defn

for n in range(150):
    psi_list = [eval(make_eigenfunction_n(n,omega)) for omega in all_omega]
    exec("psi_list_"+str(n)+" = psi_list")


# Actually make the N dimensional eigenfunction

def prod(iterable):
    return reduce(operator.mul, iterable, 1) # prod([2,3,4,5]) = factorial(5)

def Psi(t,*r):
    return psi_list_4[0](t,r[0])*psi_list_2[1](t,r[1])

#print(Psi(0.0,[1.0,1.0]))

if __name__ == "__main__":
    print("\nRunning Diagnostic Tests...\n")
    
    print("Potential Strength:")
    print("\omega = "+str(all_omega))
    
    print("\nCheck that the normalisation of the exported wavefunction is 1.0 (2D only for now...):")
    
    times = 10
    comb = [i/float(times)/2.0 for i in range(-20*times,20*times)]
    delta_comb = comb[7]-comb[6]
    sum = 0
    count = 0
    for x in comb: # metacode the loop for N dimensions?
        for y in comb:
            a = Psi(0.0,[x,y])
            sum += abs(a)**2 * delta_comb**2
            count += 1
        print count
    
    print(str.format('{0:.12f}',sum.real))